#ifndef D2_THREAD_POOL_HPP
#define D2_THREAD_POOL_HPP

#include <cassert>

#include <condition_variable>
#include <functional>
#include <future>
#include <memory>
#include <mutex>
#include <thread>
#include <type_traits>
#include <utility>
#include <vector>
#include <queue>
#include <algorithm>

namespace d2 {

namespace detail {

#if __cplusplus >= 202300L    
    using move_only_task_function = std::move_only_function<void()>;
#else
class move_only_task_function
{
    struct callable
    {
        virtual ~callable() = default;
        virtual void invoke() = 0;
    };

    template <typename F>
    struct callable_impl final : callable
    {
        F f_;
        explicit callable_impl(F&& f) : f_(std::move(f)) {}
        ~callable_impl() override = default;

        void invoke() override { f_(); }
    };

    std::unique_ptr<callable> impl; 

public:
    move_only_task_function() noexcept = default;
    move_only_task_function(std::nullptr_t) noexcept {}
    move_only_task_function(move_only_task_function const&) = delete;
    move_only_task_function(move_only_task_function&&) noexcept = default;

    template<typename F, 
            typename = std::enable_if_t<!std::is_same_v<std::decay_t<F>, move_only_task_function> 
                && std::is_invocable_v<F>
                && std::is_move_constructible_v<F> && !std::is_copy_constructible_v<F>
            >>
    move_only_task_function(F&& f) 
        : impl(std::make_unique<callable_impl<std::decay_t<F>>>(std::forward<F>(f))) 
    {
    }
    
    move_only_task_function& operator = (move_only_task_function const&) = delete;
    move_only_task_function& operator = (move_only_task_function&&) noexcept = default;
    move_only_task_function& operator = (std::nullptr_t) noexcept { impl.reset(); return *this; }

    void swap(move_only_task_function& other) noexcept { impl.swap(other.impl); }
    
    explicit operator bool() const noexcept { return impl != nullptr; }
    
    void operator () () const
    {
        if (!impl) { throw std::bad_function_call(); }        
        impl->invoke();
    }
};
#endif

//------------------------------
} // namespace detail 

#ifdef USE_MOVE_ONLY_FUNCTION_FOR_INTERNAL_TASK
    using task_function = detail::move_only_task_function;
#else
    using task_function = std::function<void()>;
#endif

// 线程池停止后（或停止过程）中仍然提交任务
struct TaskSubmittedAfterStopped : std::runtime_error
{
    TaskSubmittedAfterStopped(char const* err)    
        : std::runtime_error(err)
    {        
    }

    ~TaskSubmittedAfterStopped() noexcept override {}
};

class ThreadPool final
{
public:
    explicit ThreadPool(size_t numThreads = std::thread::hardware_concurrency())
        : stop(false), tasks_counter(numThreads)
    {
        std::for_each(tasks_counter.begin(), tasks_counter.end(), [](auto& a) {a = 0;});

        for (size_t i = 0; i < numThreads; ++i)
        {
            workers.emplace_back([this, i] () -> void
            {
                task_function task;

                for(;;)
                {
                    {
                        std::unique_lock<std::mutex> lock(this->queueMutex);
                        this->condition.wait(lock, [this] { return this->stop || !this->tasks.empty(); });

                        if (this->stop && this->tasks.empty())
                        {
                            break;
                        }

                        // stop 之后的任务，仍然要执行
                        // 下次 condition.wait ，将因 this->stop 为真而不实质等待
                        task = std::move(this->tasks.front());
                        this->tasks.pop();
                    }
                    
                    task(); // 不需要包装 try/cache ，因为 task 是 std::packate_task<R>
                    ++tasks_counter[i]; // 任务完成的再计数（因为记录的是完成数）
                }
            });
        }
    }

    ~ThreadPool() noexcept
    {
        try
        {
            Stop();
        }
        catch (...)
        {
        }
    }

    size_t GetSize() const noexcept
    {
        return workers.size();
    }

    // 取指定线程完成执行的任务数
    std::size_t GetCompletedTaskCount(std::size_t index) const noexcept
    {
        if (index >= tasks_counter.size()) [[unlikely]]
        {
            return 0;
        }

        return tasks_counter[index];        
    }

    // 取等待执行的任务总数
    size_t GetPendingTaskCount() const noexcept
    {
        std::scoped_lock lock (this->queueMutex);
        return tasks.size();
    }

    template<typename F, typename... Args>
    auto Submit(F&& f, Args&&... args)
    {
        using ReturnType = typename std::invoke_result_t<F, Args...>;
        using PackagedTaskType = std::packaged_task<ReturnType()>;

        #ifdef USE_MOVE_ONLY_FUNCTION_FOR_INTERNAL_TASK
            auto task = PackagedTaskType(std::bind(std::forward<F>(f), std::forward<Args>(args)...));
            auto res = task.get_future();
        #else
            auto task = std::make_shared<PackagedTaskType>(std::bind(std::forward<F>(f), std::forward<Args>(args)...));
            auto res = task->get_future();
        #endif

        {
            std::scoped_lock lock(queueMutex);
            
            if (!stop)
            {
                #ifdef USE_MOVE_ONLY_FUNCTION_FOR_INTERNAL_TASK
                tasks.emplace([task = std::move(task)] () mutable { task(); });
                #else
                tasks.emplace([task] () { (*task)(); });
                #endif
            }
            else
            {
                throw TaskSubmittedAfterStopped("ThreadPool::Submit() called after Stop()");
            }
        }

        condition.notify_one(); //只会唤醒一个线程
        return res;
    }

    void Stop() noexcept(false)
    {
        {
            std::scoped_lock lock(queueMutex);
            
            if (stop)
            {
                return; // 如果已经 stop 过，直接返回
            }

            stop = true;
        }

        condition.notify_all();
        
        for (std::thread &worker : workers)
        {
            if (worker.joinable())
            {
                worker.join();
            }
        }
    }

private:
    bool stop;

    std::vector<std::thread> workers;
    std::queue<task_function> tasks;
    
    std::mutex mutable queueMutex;
    std::condition_variable condition;

private:
    std::vector<std::atomic_size_t> tasks_counter;
};

} // namespace d2

#endif // D2_THREAD_POOL_HPP